import torch
import matplotlib.pyplot as plt
import numpy as np
import torch.utils.data as data
import config
import torchvision
from net import YoloV1Net
from dataset import YoloV1Dataset
from loss import YoloV1Loss

print ("Loading the datasets")
data_tv = YoloV1Dataset ("./datasets/train.csv", "./datasets/images", "./datasets/labels", normalize=True)
train_data, val_data = data.random_split(data_tv, [0.8, 0.2])

train_loader = data.DataLoader (
    dataset=train_data,
    batch_size=config.BATCH_SIZE,
    shuffle=True
)
val_loader = data.DataLoader (
    dataset=val_data,
    batch_size=config.BATCH_SIZE,
    shuffle=True
)

model = YoloV1Net ()
model.to ('cuda')

print (model)

optimizer = torch.optim.Adam (model.parameters(), lr=config.LR)
loss_func = YoloV1Loss ()

print (f"Epochs : {config.EPOCHS}\nTrain_set Length : {train_data.__len__ ()}\nBatch Size : {config.BATCH_SIZE}")
print ("Start Training ...")
for epoch in range (config.EPOCHS) :
    train_loss = []
    for step, (input, label) in enumerate (train_loader) :
        input, label = input.to ('cuda'), label.to ('cuda')
        out = model (input)
        loss = loss_func (out, label)

        optimizer.zero_grad ()
        loss.backward ()
        optimizer.step()

        if step % 10 == 0 :
            print (f"Epoch : {epoch}\tLoss : {loss.item ()}")
            train_loss.append(loss.item ())

plt.plot (train_loss)
plt.show ()
torch.save(model, "./model.pt")
